import os

import numpy as np
from osgeo import gdal, gdalconst
import csv

'''
草地理论承载力是指在一定时间内，特定草地生态系统能够持续提供的牧草
资源所能支持的最大家畜或野生动物数量，而不损害草地生态系统的生产力和可
持续性。评估草地可以承受的最大蓄牧生产量，以确定其长期可持续性。以下为
计算公式和参数参考值：
    AFYi=NPP×EF×UE
    TSCC=(AFY×GRi)/(A×365)
其中：AFY;为第i种草地类型的可利用牧草产量，NPP为地上生产力，EF为草地
标准干草折算系数，UE为草地合理利用率。TSCC为某种草地类型的合理载畜量，
GR;为第i种草地放牧采食率，A为标准羊单位日食量(1.8 kg/d),365为放牧
天数。干草折算系数和利用率见表2。输入数据为2.2.1生成的NPP生产力栅格、
草地承载力关键参数表格CSV(包含不同草地类型的放牧采食率GRi、干草折算
系数EF、标准羊单位日食量A、草地合理利用率UE),输出数据为TSCC草地承载
力栅格数据。

'''


def read_type_values_from_csv(csv_path):
    """
    从 CSV 文件中读取草地类型与 EF、UE、GR 的映射。

    参数:
    csv_path (str): CSV 文件的路径。

    返回:
    dict: 草地类型与 EF、UE、GR 的映射字典，格式为 {类型: (EF, UE, GR)}。
    """
    type_values = {}

    with open(csv_path, mode='r', newline='') as file:
        reader = csv.reader(file)
        headers = next(reader)  # 读取表头
        if headers != ["type", "EF", "UE", "GR"]:
            raise ValueError("CSV 文件列标题必须为 'type', 'EF', 'UE', 'GR'")
        for row in reader:
            grass_type = int(row[0])  # 草地类型
            ef = float(row[1])  # EF 值
            ue = float(row[2])  # UE 值
            gr = float(row[3])  # GR 值
            type_values[grass_type] = (ef, ue, gr)
    return type_values


def get_unique_types_from_raster(raster_path):
    """
    从 TIF 文件中获取唯一的草地类型值。

    参数:
    raster_path (str): 草地类型图的路径（GeoTIFF 文件）。

    返回:
    set: 唯一的草地类型值集合。
    """
    dataset = gdal.Open(raster_path, gdalconst.GA_ReadOnly)
    if dataset is None:
        raise FileNotFoundError(f"无法打开文件: {raster_path}")

    # 读取第一个波段的数据
    band = dataset.GetRasterBand(1)
    data = band.ReadAsArray()

    # 获取 NoData 值
    nodata_value = band.GetNoDataValue()

    # 获取唯一的草地类型值（忽略 NoData 值）
    unique_types = set(np.unique(data))
    if nodata_value is not None:
        unique_types.discard(nodata_value)  # 移除 NoData 值

    dataset = None
    return unique_types


def resample_raster_to_memory(input_path, reference_dataset):
    """
    将栅格图像重采样到参考栅格的尺寸和空间坐标，并返回内存中的栅格数据。

    参数:
    input_path (str): 输入栅格文件的路径。
    reference_dataset: 参考数据集（用于获取目标尺寸和空间坐标）。

    返回:
    numpy.ndarray: 重采样后的栅格数据。
    """
    # 获取参考栅格的尺寸和空间坐标
    target_width = reference_dataset.RasterXSize
    target_height = reference_dataset.RasterYSize
    geo_transform = reference_dataset.GetGeoTransform()
    projection = reference_dataset.GetProjection()

    # 创建内存中的输出数据集
    driver = gdal.GetDriverByName("MEM")
    output_dataset = driver.Create("", target_width, target_height, 1, gdalconst.GDT_Float32)

    # 设置输出数据集的地理变换和投影
    output_dataset.SetGeoTransform(geo_transform)
    output_dataset.SetProjection(projection)

    # 使用 gdal.Warp 进行重采样
    gdal.Warp(output_dataset, input_path, width=target_width, height=target_height, resampleAlg="bilinear")

    # 读取重采样后的数据
    band = output_dataset.GetRasterBand(1)
    data = band.ReadAsArray()

    # 关闭数据集
    output_dataset = None

    return data


def multiply_grass_type_raster(input_path, npp_path, output_path, type_values, save_npp=False):
    """
    将草地类型图与 EF、UE、GR 以及 NPP 相乘，并保存结果。

    参数:
    input_path (str): 输入草地类型图的路径（GeoTIFF 文件）。
    npp_path (str): 输入 NPP 栅格图的路径（GeoTIFF 文件）。
    output_path (str): 输出文件的路径。
    type_values (dict): 草地类型与 EF、UE、GR 的映射字典，格式为 {类型: (EF, UE, GR)}。
    save_npp (bool): 是否保存重采样后的 NPP 数据。
    """
    # 打开输入草地类型图
    dataset = gdal.Open(input_path, gdalconst.GA_ReadOnly)
    if dataset is None:
        raise FileNotFoundError(f"无法打开文件: {input_path}")

    # 读取第一个波段的数据
    band = dataset.GetRasterBand(1)
    data = band.ReadAsArray()

    # 获取 NoData 值
    nodata_value = band.GetNoDataValue()

    # 获取草地类型图的尺寸
    rows, cols = data.shape

    # 打开 NPP 栅格图
    npp_dataset = gdal.Open(npp_path, gdalconst.GA_ReadOnly)
    if npp_dataset is None:
        raise FileNotFoundError(f"无法打开 NPP 文件: {npp_path}")

    # 检查 NPP 图的尺寸是否与草地类型图一致
    npp_band = npp_dataset.GetRasterBand(1)
    npp_data = npp_band.ReadAsArray()
    if npp_data.shape != (rows, cols):
        print("NPP 图尺寸与草地类型图不一致，正在重采样...")
        npp_data = resample_raster_to_memory(npp_path, dataset)  # 重采样到内存中

        # 如果 save_npp 为 True，保存重采样后的 NPP 数据
        if save_npp:
            npp_output_path = "resampled_npp.tif"
            save_raster(npp_output_path, npp_data, dataset)
            print(f"重采样后的 NPP 数据已保存到: {npp_output_path}")

    # 初始化输出数据
    result = np.zeros_like(data, dtype=np.float32)

    # 遍历草地类型，计算乘积
    for grass_type, (ef, ue, gr) in type_values.items():
        mask = (data == grass_type)  # 找到当前类型的区域
        result[mask] = (npp_data[mask] * ef * ue) * gr / (1.8 * 365)

    # 将 NoData 区域设置为 NoData 值
    if nodata_value is not None:
        result[data == nodata_value] = nodata_value

    # 保存结果
    save_raster(output_path, result, dataset)
    print(f"结果已保存到: {output_path}")

    # 关闭数据集
    dataset = None
    npp_dataset = None


def save_raster(output_path, data, reference_dataset):
    """
    将栅格数据保存为 GeoTIFF 文件。

    参数:
    output_path (str): 输出文件的路径。
    data (numpy.ndarray): 要保存的栅格数据。
    reference_dataset: 参考数据集（用于获取地理信息和投影）。
    """
    driver = gdal.GetDriverByName("GTiff")
    rows, cols = data.shape
    out_dataset = driver.Create(output_path, cols, rows, 1, gdalconst.GDT_Float32)

    # 设置地理变换和投影
    out_dataset.SetGeoTransform(reference_dataset.GetGeoTransform())
    out_dataset.SetProjection(reference_dataset.GetProjection())

    # 写入数据
    out_band = out_dataset.GetRasterBand(1)
    out_band.WriteArray(data)

    # 关闭数据集
    out_band.FlushCache()
    out_dataset = None


if __name__ == "__main__":
    # CSV 文件路径
    csv_path = "F:/test/参数-test.csv"

    # 输入草地类型图路径
    input_path = "F:/test/lu2022.tif"

    # 输入 NPP 栅格图路径
    npp_path = "F:/test/多样性/NPP2022.tif"

    # 输出文件路径
    output_path = "F:/test/grass_type_product_with_npp.tif"

    # 是否保存重采样后的 NPP 数据
    save_npp = True  # 设置为 True 保存，False 不保存

    # 从 CSV 文件中读取草地类型与 EF、UE、GR 的映射
    try:
        type_values = read_type_values_from_csv(csv_path)
        print("成功读取 CSV 文件！")
    except Exception as e:
        print(f"读取 CSV 文件时发生错误: {e}")
        type_values = {}  # 如果读取失败，使用空字典

    # 获取 TIF 文件中的唯一草地类型值
    try:
        unique_types = get_unique_types_from_raster(input_path)
        print(f"TIF 文件中的草地类型值: {unique_types}")
    except Exception as e:
        print(f"读取 TIF 文件时发生错误: {e}")
        unique_types = set()

    # 检查 CSV 文件中的类型数目是否与 TIF 文件中的类型数目一致
    if type_values and unique_types:
        csv_types = set(type_values.keys())
        if csv_types != unique_types:
            raise ValueError(f"CSV 文件中的类型数目与 TIF 文件中的类型数目不一致\n"
                             f"CSV 文件中的类型: {csv_types}\n"
                             f"TIF 文件中的类型: {unique_types}")
        else:
            print("CSV 文件中的类型数目与 TIF 文件中的类型数目一致！")

    # 调用函数
    if type_values:
        try:
            multiply_grass_type_raster(input_path, npp_path, output_path, type_values, save_npp=save_npp)
        except Exception as e:
            print(f"处理栅格数据时发生错误: {e}")
